package com.geek.gkbgw.framework.trace;


import com.geek.gkbgw.framework.trace.model.TraceInfo;
import org.slf4j.MDC;

/**
 * <h1>全局ID本地缓存封装对象。</h1><br/>
 *
 * <h1><br/><a href="https://u.geekbang.org/subject/arch3/100614201">Java 业务架构实战营：https://u.geekbang.org/subject/arch3/100614201</a></h1><br/><h1><a href="https://gitee.com/ylimhhmily/BizArchPracticeCamp">Java 业务架构实战营代码地址：https://gitee.com/ylimhhmily/BizArchPracticeCamp</a></h1><br/><h1><a href="https://u.geekbang.org/subject/java4th/1001148?source=app_share">高级 Java 工程体系课 2023 版：https://u.geekbang.org/subject/java4th/1001148?source=app_share</a></h1><br/><h1><a href="https://time.geekbang.org/column/intro/100312101?tab=catalog">Dubbo 源码剖析与实战：https://time.geekbang.org/column/intro/100312101?tab=catalog</a></h1>
 *
 * @author hmilyylimh
 *         ^_^
 * @version 0.0.1
 *         ^_^
 * @date 2022-02-19
 *
 */
public class TraceContext {

    private static final ThreadLocal<TraceInfo> CURRENT = new ThreadLocal<TraceInfo>();

    public static TraceInfo getOrCreate() {
        if (getCurrentTrace() == null) {
            setCurrentSpan(TraceInfo.create());
        }

        return getCurrentTrace();
    }

    public static TraceInfo getCurrentTrace() {
        return CURRENT.get();
    }

    public static void setCurrentSpan(TraceInfo traceInfo) {
        if (traceInfo != null && getCurrentTrace() == null) {
            CURRENT.set(traceInfo);
            fillInMdc();
        }
    }

    /**
     * <h2>获取 traceId 值。</h2>
     *
     * @return
     */
    public static String getTraceId() {
        TraceInfo currentTrace = getCurrentTrace();
        return currentTrace == null ? "" : currentTrace.getTraceId();
    }

    /**
     * <h2>mark相同则清楚，否则存在超5分钟也干掉。</h2>
     *
     * @param mark
     */
    public static void removeCurrentTrace(String mark) {
        TraceInfo traceInfo = getCurrentTrace();
        if (traceInfo == null) {
            return;
        }

        String currMark = traceInfo.getMarker();
        if (mark.equals(currMark) || traceInfo.isPassed()) {
            clearMdc();
            CURRENT.remove();
        }
    }

    private static void fillInMdc() {
        TraceInfo traceInfo = getCurrentTrace();
        if (traceInfo == null) {
            return;
        }

        MDC.put(TraceInfo.TRACE_ID_NAME, traceInfo.getTraceId());
        MDC.put(TraceInfo.SPAN_ID_NAME, traceInfo.getSpanId());
    }

    private static void clearMdc() {
        MDC.remove(TraceInfo.TRACE_ID_NAME);
        MDC.remove(TraceInfo.SPAN_ID_NAME);
    }
}